import rosbag
import cv2
import numpy as np
import os
import sys

bag = rosbag.Bag(sys.argv[1])

if not os.path.exists("imu0"):
    os.makedirs("imu0")
imu_file = open('imu0/data.txt', 'w')
imu_file.write('#timestamp [ns],w_RS_S_x [rad s^-1],w_RS_S_y [rad s^-1],w_RS_S_z [rad s^-1],a_RS_S_x [m s^-2],a_RS_S_y [m s^-2],a_RS_S_z [m s^-2]\n')

cam_folder = 'cam0/data'
if not os.path.exists(cam_folder):
    os.makedirs(cam_folder)
cam_file = open('cam0/data.txt', 'w')

if not os.path.exists("encoder0"):
    os.makedirs("encoder0")
encoder_file = open('encoder0/data.csv', 'w')
encoder_file.write('timestamp,encoder_data\n')

for topic, msg, t in bag.read_messages():
    if topic == '/wit/imu':
        imu_file.write(str(msg.header.stamp.to_nsec()) + ',' + str(msg.angular_velocity.x) + ',' + str(msg.angular_velocity.y) + ',' + str(msg.angular_velocity.z) + ',' + str(msg.linear_acceleration.x) + ',' + str(msg.linear_acceleration.y) + ',' + str(msg.linear_acceleration.z) + '\n')
    elif topic == '/usb_cam/image_raw':
        img = np.fromstring(msg.data, dtype=np.uint8)
        img = img.reshape((msg.height, msg.width, 3))
        img_name = str(msg.header.stamp.to_nsec()) + '.png'
        cv2.imwrite(os.path.join(cam_folder, img_name), img)
        cam_file.write(str(msg.header.stamp.to_nsec()) + '\n')
    elif topic == '/encoder_data':
        encoder_file.write(str(msg.stamp.to_nsec()) + ',' + msg.frame_id + '\n')

imu_file.close()
cam_file.close()
encoder_file.close()

bag.close()
